import torch

# ==== 定义 Poisson 噪声类 ====
class PoissonNoise(torch.nn.Module):
    """Poisson 噪声 (依赖像素强度)"""
    def __init__(self, scale=255.0):
        super().__init__()
        self.scale = scale

    def forward(self, x):
        # x: [B,C,H,W] in [0,1]
        noisy = torch.poisson(x * self.scale) / self.scale
        return torch.clamp(noisy, 0, 1)